import json
import mysql.connector
import csv
from datetime import datetime
from typing import Dict, Any, List


class TelegramMessageProcessor:
    def __init__(self, db_config: Dict[str, str], csv_file_path: str = None, qa_csv_file_path: str = None):
        self.db_config = db_config
        self.csv_file_path = csv_file_path
        self.qa_csv_file_path = qa_csv_file_path
        self.conn = None
        self.cursor = None
        self.csv_file = None
        self.csv_writer = None
        self.qa_csv_file = None
        self.qa_csv_writer = None

        # 定义需要过滤的指令列表
        self.commands_to_filter = [
            '/payin', '/budan', '/bd', '/piq', '/poq', '/payout', '/proof',
            '/tx', '/cbi', '/cbo', '/qu', '/pri', '/price', '/u', '/xiafa',
            '/bal', '/help', '/rate', '/bind', 'ok', '稍等','订单处理失败','订单已成功处理','尊敬的商户',
            '订单已完结','补单处理中','password','/doc','测试单，代收，回调：成功状态，谢谢'
        ]

        # 存储指令消息的ID，用于过滤回复这些指令的回复消息
        self.command_message_ids = set()

        # 存储所有消息的映射，用于查找回复的消息
        self.all_messages = {}

    def connect_db(self):
        """连接到MySQL数据库"""
        try:
            self.conn = mysql.connector.connect(**self.db_config)
            self.cursor = self.conn.cursor()
            print("成功连接到数据库")
            # 连接成功后检查并创建表（如果不存在）
            self.create_tables()
        except mysql.connector.Error as err:
            print(f"数据库连接错误: {err}")
            raise

    def create_tables(self):
        """创建必要的数据库表（如果不存在）"""
        try:
            # 创建messages表
            create_messages_table = """
            CREATE TABLE IF NOT EXISTS messages (
                id BIGINT PRIMARY KEY,
                from_user VARCHAR(255),
                from_id VARCHAR(100),
                text_content TEXT,
                reply_to_message_id BIGINT,
                message_type VARCHAR(50),
                date DATETIME,
                date_unixtime BIGINT,
                edited_date DATETIME,
                edited_unixtime BIGINT
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
            """

            # 创建reactions表
            create_reactions_table = """
            CREATE TABLE IF NOT EXISTS reactions (
                id INT AUTO_INCREMENT PRIMARY KEY,
                message_id BIGINT,
                emoji VARCHAR(50),
                reaction_count INT,
                FOREIGN KEY (message_id) REFERENCES messages(id) ON DELETE CASCADE
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
            """

            # 创建reaction_users表
            create_reaction_users_table = """
            CREATE TABLE IF NOT EXISTS reaction_users (
                id INT AUTO_INCREMENT PRIMARY KEY,
                reaction_id INT,
                user_name VARCHAR(255),
                user_id VARCHAR(100),
                reaction_date DATETIME,
                FOREIGN KEY (reaction_id) REFERENCES reactions(id) ON DELETE CASCADE
            ) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4;
            """

            # 执行创建表的SQL语句
            self.cursor.execute(create_messages_table)
            self.cursor.execute(create_reactions_table)
            self.cursor.execute(create_reaction_users_table)
            self.conn.commit()
            print("数据库表检查完成（已存在的表不会重复创建）")

        except mysql.connector.Error as err:
            print(f"创建表错误: {err}")
            self.conn.rollback()
            raise

    def setup_csv_writer(self):
        """设置CSV写入器"""
        if self.csv_file_path:
            try:
                # 使用 utf-8-sig 编码，这样 Excel 打开时不会乱码
                self.csv_file = open(self.csv_file_path, 'w', newline='', encoding='utf-8-sig')
                self.csv_writer = csv.writer(self.csv_file)
                # 写入CSV表头
                self.csv_writer.writerow([
                    'id', 'from_user', 'from_id', 'text_content', 'reply_to_message_id',
                    'message_type', 'date', 'date_unixtime', 'edited_date', 'edited_unixtime'
                ])
                print(f"CSV文件已创建: {self.csv_file_path}")
            except Exception as e:
                print(f"创建CSV文件错误: {e}")
                raise

        if self.qa_csv_file_path:
            try:
                # 创建问答对CSV文件
                self.qa_csv_file = open(self.qa_csv_file_path, 'w', newline='', encoding='utf-8-sig')
                self.qa_csv_writer = csv.writer(self.qa_csv_file)
                # 写入问答对CSV表头
                self.qa_csv_writer.writerow([
                    'question_id', 'question_user', 'question_text',
                    'answer_id', 'answer_user', 'answer_text',
                    'date'
                ])
                print(f"问答对CSV文件已创建: {self.qa_csv_file_path}")
            except Exception as e:
                print(f"创建问答对CSV文件错误: {e}")
                raise

    def disconnect_db(self):
        """断开数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()
        if self.csv_file:
            self.csv_file.close()
        if self.qa_csv_file:
            self.qa_csv_file.close()
        print("数据库连接和CSV文件已关闭")

    def parse_text_entities(self, text_entities: List[Any]) -> str:
        """解析消息文本实体，提取纯文本内容"""
        full_text = ""
        for entity in text_entities:
            if isinstance(entity, dict):
                if 'text' in entity:
                    full_text += entity['text']
            elif isinstance(entity, str):
                full_text += entity
        return full_text

    def is_command_message(self, text: str) -> bool:
        """检查消息是否以需要过滤的指令开头"""
        if not text:
            return False

        # 去除首尾空白字符
        stripped_text = text.strip()

        # 检查是否以任何指令开头
        for command in self.commands_to_filter:
            if stripped_text.startswith(command):
                return True

        return False

    def write_to_csv(self, message_data: Dict[str, Any]):
        """将消息数据写入CSV文件"""
        if self.csv_writer:
            try:
                self.csv_writer.writerow([
                    message_data.get('id'),
                    message_data.get('from_user'),
                    message_data.get('from_id'),
                    message_data.get('text_content'),
                    message_data.get('reply_to_message_id'),
                    message_data.get('message_type'),
                    message_data.get('date'),
                    message_data.get('date_unixtime'),
                    message_data.get('edited_date'),
                    message_data.get('edited_unixtime')
                ])
            except Exception as e:
                print(f"写入CSV错误: {e}")

    def write_qa_to_csv(self, question_data: Dict[str, Any], answer_data: Dict[str, Any]):
        """将问答对数据写入CSV文件"""
        if self.qa_csv_writer:
            try:
                self.qa_csv_writer.writerow([
                    question_data.get('id'),
                    question_data.get('from_user'),
                    question_data.get('text_content'),
                    answer_data.get('id'),
                    answer_data.get('from_user'),
                    answer_data.get('text_content'),
                    answer_data.get('date')
                ])
            except Exception as e:
                print(f"写入问答对CSV错误: {e}")

    def process_message(self, message: Dict[str, Any]):
        """处理单条消息并插入数据库和CSV"""
        # 只处理文本消息
        if message.get('type') != 'message':
            return

        # 提取消息内容
        message_id = message.get('id')
        message_type = message.get('type')
        date_str = message.get('date')
        date_unixtime = message.get('date_unixtime')
        from_user = message.get('from')  # 提取消息发送者
        from_id = message.get('from_id')
        reply_to_message_id = message.get('reply_to_message_id')
        edited_date = message.get('edited')
        edited_unixtime = message.get('edited_unixtime')

        # 过滤“小深助手-Pakistan”的消息
        if from_user == '小深助手-Pakistan':
            print(f"跳过小深助手-Pakistan消息: ID={message_id}")
            return

        # 处理文本内容
        text_entities = message.get('text_entities', [])
        text_content = self.parse_text_entities(text_entities)

        # 如果没有text_entities，尝试使用text字段
        if not text_content and 'text' in message:
            if isinstance(message['text'], list):
                text_content = self.parse_text_entities(message['text'])
            else:
                text_content = str(message['text'])

        # 检查是否为指令消息，如果是则记录ID并跳过
        if self.is_command_message(text_content):
            print(f"跳过指令消息: {text_content[:50]}...")
            self.command_message_ids.add(message_id)
            return

        # 检查是否为回复指令消息的回复消息，如果是则跳过
        if reply_to_message_id and reply_to_message_id in self.command_message_ids:
            print(f"跳过回复指令的消息: {text_content[:50]}...")
            return

        # 转换日期格式
        date_obj = datetime.fromisoformat(date_str) if date_str else None
        edited_date_obj = datetime.fromisoformat(edited_date) if edited_date else None

        # 准备消息数据
        message_data = {
            'id': message_id,
            'from_user': from_user,
            'from_id': from_id,
            'text_content': text_content,
            'reply_to_message_id': reply_to_message_id,
            'message_type': message_type,
            'date': date_obj,
            'date_unixtime': date_unixtime,
            'edited_date': edited_date_obj,
            'edited_unixtime': edited_unixtime
        }

        # 存储消息到映射中，用于后续查找
        self.all_messages[message_id] = message_data

        # 检查是否为回复消息，如果是则创建问答对
        if reply_to_message_id and reply_to_message_id in self.all_messages:
            question_data = self.all_messages[reply_to_message_id]
            self.write_qa_to_csv(question_data, message_data)

        # 插入消息到数据库
        try:
            insert_query = """
            INSERT INTO messages (id, from_user, from_id, text_content, reply_to_message_id, 
                                message_type, date, date_unixtime, edited_date, edited_unixtime)
            VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            self.cursor.execute(insert_query, (
                message_id, from_user, from_id, text_content, reply_to_message_id,
                message_type, date_obj, date_unixtime, edited_date_obj, edited_unixtime
            ))

            # 同时写入CSV文件
            self.write_to_csv(message_data)

            # 处理反应(reactions)
            if 'reactions' in message:
                self.process_reactions(message_id, message['reactions'])

        except mysql.connector.Error as err:
            print(f"插入消息错误: {err}")

    def process_reactions(self, message_id: int, reactions: List[Dict[str, Any]]):
        """处理消息的反应"""
        for reaction in reactions:
            emoji = reaction.get('emoji')
            count = reaction.get('count')

            # 插入反应
            reaction_query = """
            INSERT INTO reactions (message_id, emoji, reaction_count)
            VALUES (%s, %s, %s)
            """
            self.cursor.execute(reaction_query, (message_id, emoji, count))
            reaction_db_id = self.cursor.lastrowid

            # 处理反应的用户
            if 'recent' in reaction:
                for user in reaction['recent']:
                    user_name = user.get('from')
                    user_id = user.get('from_id')
                    reaction_date_str = user.get('date')
                    reaction_date = datetime.fromisoformat(reaction_date_str) if reaction_date_str else None

                    user_query = """
                    INSERT INTO reaction_users (reaction_id, user_name, user_id, reaction_date)
                    VALUES (%s, %s, %s, %s)
                    """
                    self.cursor.execute(user_query, (reaction_db_id, user_name, user_id, reaction_date))

    def process_json_file(self, file_path: str):
        """处理JSON文件"""
        try:
            with open(file_path, 'r', encoding='utf-8') as file:
                data = json.load(file)

            messages = data.get('messages', [])
            print(f"找到 {len(messages)} 条消息")

            filtered_count = 0
            processed_count = 0

            # 第一遍：识别所有指令消息并记录它们的ID
            for message in messages:
                if message.get('type') == 'message':
                    text_entities = message.get('text_entities', [])
                    text_content = self.parse_text_entities(text_entities)

                    if not text_content and 'text' in message:
                        if isinstance(message['text'], list):
                            text_content = self.parse_text_entities(message['text'])
                        else:
                            text_content = str(message['text'])

                    if self.is_command_message(text_content):
                        message_id = message.get('id')
                        self.command_message_ids.add(message_id)
                        filtered_count += 1

            # 第二遍：处理消息，过滤指令消息和回复指令消息的回复
            for message in messages:
                if message.get('type') != 'message':
                    continue

                message_id = message.get('id')
                reply_to_message_id = message.get('reply_to_message_id')

                # 如果是指令消息，跳过
                if message_id in self.command_message_ids:
                    filtered_count += 1
                    continue

                # 如果是回复指令消息的回复，跳过
                if reply_to_message_id and reply_to_message_id in self.command_message_ids:
                    filtered_count += 1
                    continue

                # 检查消息本身是否为指令消息（双重检查）
                text_entities = message.get('text_entities', [])
                text_content = self.parse_text_entities(text_entities)

                if not text_content and 'text' in message:
                    if isinstance(message['text'], list):
                        text_content = self.parse_text_entities(message['text'])
                    else:
                        text_content = str(message['text'])

                if self.is_command_message(text_content):
                    filtered_count += 1
                    self.command_message_ids.add(message_id)  # 记录新发现的指令消息
                    continue

                self.process_message(message)
                processed_count += 1

            self.conn.commit()
            print(f"数据处理完成: 处理了 {processed_count} 条消息, 过滤了 {filtered_count} 条消息")

        except Exception as e:
            print(f"处理文件时出错: {e}")
            self.conn.rollback()


# 使用示例
if __name__ == "__main__":
    # 数据库配置
    db_config = {
        'host': 'localhost',
        'user': 'root',
        'password': 'root',
        'database': 'telegram_chat'
    }

    # CSV文件路径（请根据实际路径调整）
    csv_file_path = 'D:\\yuanbei\\客服问答机器人需求表\\清洗数据\\data_process\\data_process_b\\12298\\data_clean.csv'
    qa_csv_file_path = 'D:\\yuanbei\\客服问答机器人需求表\\清洗数据\\data_process\\data_process_b\\12298\\qa_pairs.csv'

    # 创建处理器实例
    processor = TelegramMessageProcessor(db_config, csv_file_path, qa_csv_file_path)

    try:
        # 连接到数据库
        processor.connect_db()

        # 设置CSV写入器
        processor.setup_csv_writer()

        # 处理JSON文件（请根据实际路径调整）
        processor.process_json_file(
            'D:\\yuanbei\\客服问答机器人需求表\\清洗数据\\data_process\\data_process_b\\12298\\result.json')

    finally:
        # 断开数据库连接
        processor.disconnect_db()
